/* -*- P4_16 -*- */
#include <core.p4>
#include <v1model.p4>

const bit<16> TYPE_IPV4 = 0x800;
const bit<16> TYPE_IPV6 = 0x86DD;
const bit<16> TYPE_ARP = 0X0806;
const bit<8> IPPROTO_ICMP = 0X01;
const bit<8> ICMP_ECHO_REQUEST = 8;
const bit<8> ICMP_ECHO_REPLY = 0;
const bit<8> ICMPV6_NS = 135;    //Neighbor Solicitation
const bit<8> ICMPV6_NA = 136;    //Neighbor Advertisement
/*************************************************************************
*********************** H E A D E R S  ***********************************
*************************************************************************/

typedef bit<9>  egressSpec_t;
typedef bit<48> macAddr_t;
typedef bit<32> ip4Addr_t;
typedef bit<128> ip6Addr_t;

header ethernet_t {
    macAddr_t dstAddr;
    macAddr_t srcAddr;
    bit<16>   etherType;
}

header ipv4_t {
    bit<4>    version;
    bit<4>    ihl;
    bit<8>    diffserv;
    bit<16>   totalLen;
    bit<16>   identification;
    bit<3>    flags;
    bit<13>   fragOffset;
    bit<8>    ttl;
    bit<8>    protocol;
    bit<16>   hdrChecksum;
    ip4Addr_t srcAddr;
    ip4Addr_t dstAddr;
}

// 新增IPv6头部定义
header ipv6_t {
    bit<4>    version;
    bit<8>    trafficClass;
    bit<20>   flowLabel;
    bit<16>   payloadLen;
    bit<8>    nextHeader;
    bit<8>    hopLimit;
    ip6Addr_t srcAddr;
    ip6Addr_t dstAddr;
}

const bit<16> ARP_OPER_REQUEST = 1;
const bit<16> ARP_OPER_REPLY =2;

//Add
header arp_t{
    bit<16> htype;  //Ethernet len
    bit<16> ptype;  //IPV4 len
    bit<8> hlen;   //MAC len
    bit<8> plen;  //IP len
    bit<16> oper;  // 1== request  2 == response
    macAddr_t sha;  //source MAC
    ip4Addr_t spa;   //source IP
    macAddr_t tha;  //destination MAC
    ip4Addr_t tpa;  //destination IP
}

header icmp_t {
    bit<8> type;
    bit<8> code;
    bit<16> checksum;
}

header icmpv6_t {
    bit<8> type;
    bit<8> code;
    bit<16> checksum;
    bit<32> reserved;
    ip6Addr_t target_addr;
}

struct metadata {
    /* empty */
    ip4Addr_t dst_ipv4;
    ip6Addr_t dst_ipv6;  // 新增：用于IPv6
    macAddr_t switch_mac;
}

struct headers {
    ethernet_t   ethernet;
    ipv4_t       ipv4;
    ipv6_t       ipv6;  // 新增IPv6头部
    arp_t        arp;
    icmp_t       icmp;
    icmpv6_t     icmpv6;
}

/*************************************************************************
*********************** P A R S E R  ***********************************
*************************************************************************/

parser MyParser(packet_in packet,
                out headers hdr,
                inout metadata meta,
                inout standard_metadata_t standard_metadata) {

    state start {
        /* TODO: add parser logic */
        transition parse_ethernet;
    }

    state parse_ethernet {
      packet.extract(hdr.ethernet);
      transition select(hdr.ethernet.etherType){
        TYPE_IPV4: parse_ipv4;
        TYPE_IPV6: parse_ipv6;
        TYPE_ARP: parse_arp;
        default: accept;
      }
    }

    state parse_arp {
          packet.extract(hdr.arp);
          meta.dst_ipv4 = hdr.arp.tpa;  // Store target protocol address in metadata
          transition accept;
      }

    state parse_ipv4{
      packet.extract(hdr.ipv4);
      transition select(hdr.ipv4.protocol) {
          IPPROTO_ICMP: parse_icmp;
          default: accept;
      }
    }

    // 新增IPv6解析状态
    state parse_ipv6 {
        packet.extract(hdr.ipv6);
        transition select(hdr.ipv6.nextHeader) {
            58: parse_icmpv6;  //58 = ICMPV6
            default: accept;
        }
    }

    state parse_icmp {
        packet.extract(hdr.icmp);
        transition accept;
    }

    state parse_icmpv6 {
        packet.extract(hdr.icmpv6);
        transition accept;
    }

  }



/*************************************************************************
************   C H E C K S U M    V E R I F I C A T I O N   *************
*************************************************************************/

control MyVerifyChecksum(inout headers hdr, inout metadata meta) {
    apply {  }
}


/*************************************************************************
**************  I N G R E S S   P R O C E S S I N G   *******************
*************************************************************************/

control MyIngress(inout headers hdr,
                  inout metadata meta,
                  inout standard_metadata_t standard_metadata) {
    action drop() {
        mark_to_drop(standard_metadata);
    }

    action ipv4_forward(macAddr_t dstAddr,egressSpec_t port) {
        /* TODO: fill out code in action body */
        standard_metadata.egress_spec = port;
        hdr.ethernet.srcAddr = hdr.ethernet.dstAddr;
        hdr.ethernet.dstAddr = dstAddr;
        hdr.ipv4.ttl = hdr.ipv4.ttl-1;

    }

    // 新增IPv6转发动作
    action ipv6_forward(macAddr_t dstAddr, egressSpec_t port) {
        standard_metadata.egress_spec = port;
        hdr.ethernet.srcAddr = hdr.ethernet.dstAddr;
        hdr.ethernet.dstAddr = dstAddr;
        hdr.ipv6.hopLimit = hdr.ipv6.hopLimit - 1;
    }

/* Add ARP reply */
    action send_arp_reply(macAddr_t macAddr) {
        //Constructing the ARP response header
        hdr.ethernet.dstAddr = hdr.arp.sha;
        hdr.ethernet.srcAddr = macAddr;

        hdr.arp.oper = ARP_OPER_REPLY;

        hdr.arp.tha = hdr.arp.sha;  // Set target hardware address to the sender's MAC
        hdr.arp.tpa = hdr.arp.spa;  // Set target protocol address to the sender's IP
        hdr.arp.sha = macAddr;
        hdr.arp.spa = meta.dst_ipv4;

        //Specify return from original pase
        standard_metadata.egress_spec = standard_metadata.ingress_port;

    }

    action send_icmp_reply() {
        macAddr_t tmp_mac;
        ip4Addr_t tmp_ip;

        //update the source IP to the current IP
        tmp_mac               = hdr.ethernet.dstAddr;
        hdr.ethernet.dstAddr  = hdr.ethernet.srcAddr;
        hdr.ethernet.srcAddr  = tmp_mac;

        tmp_ip                = hdr.ipv4.dstAddr;
        hdr.ipv4.dstAddr      = hdr.ipv4.srcAddr;
        hdr.ipv4.srcAddr      = tmp_ip;

        hdr.icmp.type         = ICMP_ECHO_REPLY;
        hdr.icmp.checksum     = 0;

        standard_metadata.egress_spec = standard_metadata.ingress_port;
    }

    action send_na(macAddr_t my_mac, ip6Addr_t my_ipv6) {
        hdr.ethernet.dstAddr = hdr.ethernet.srcAddr;
        hdr.ethernet.srcAddr = my_mac;
        hdr.ipv6.dstAddr = hdr.ipv6.srcAddr;
        hdr.ipv6.srcAddr = my_ipv6;
        hdr.icmpv6.type = 136;

        standard_metadata.egress_spec = standard_metadata.ingress_port;
    }


    table ipv4_lpm {
        key = {
            hdr.ipv4.dstAddr: lpm;
        }
        actions = {
            ipv4_forward;
            drop;
            NoAction;
        }
        size = 1024;
        default_action = NoAction();
    }

    // 新增IPv6路由表
    table ipv6_lpm {
        key = {
            hdr.ipv6.dstAddr: lpm;
        }
        actions = {
            ipv6_forward;
            drop;
            NoAction;
        }
        size = 1024;
        default_action = drop();
    }

/* Add ARP response */
    table arp_match {
        key = {
            hdr.arp.oper  : exact; // Match ARP operation (request or reply)
            hdr.arp.tpa   : lpm; // Match target protocol address using longest prefix match
        }
        actions = {
            send_arp_reply;
            drop;
        }
        const default_action = drop();
    }

    table ndp_match {
        key = {
            hdr.icmpv6.type: exact;
            hdr.icmpv6.target_addr: exact;
        }
        actions = {
            send_na;
            drop;
        }
        size = 16;
        default_action = drop();
    }

    apply {
         if(hdr.ipv4.isValid() && hdr.icmp.isValid() &&hdr.ipv4.protocol == IPPROTO_ICMP ) {
            send_icmp_reply();
         }
         else if (hdr.arp.isValid() && hdr.arp.oper == ARP_OPER_REQUEST) {
            arp_match.apply();
         }
         else if (hdr.ipv4.isValid()){
            ipv4_lpm.apply();
         }
         else if (hdr.ipv6.isValid()) {  // 新增IPv6处理
            ipv6_lpm.apply();
         }
         else if (hdr.icmpv6.isValid() && hdr.icmpv6.type == 135) {
            ndp_match.apply();
         }
    }
}

/*************************************************************************
****************  E G R E S S   P R O C E S S I N G   *******************
*************************************************************************/

control MyEgress(inout headers hdr,
                 inout metadata meta,
                 inout standard_metadata_t standard_metadata) {
    apply {  }
}

/*************************************************************************
*************   C H E C K S U M    C O M P U T A T I O N   **************
*************************************************************************/

control MyComputeChecksum(inout headers hdr, inout metadata meta) {
     apply {
        update_checksum(
            hdr.ipv4.isValid(),
            { hdr.ipv4.version,
              hdr.ipv4.ihl,
              hdr.ipv4.diffserv,
              hdr.ipv4.totalLen,
              hdr.ipv4.identification,
              hdr.ipv4.flags,
              hdr.ipv4.fragOffset,
              hdr.ipv4.ttl,
              hdr.ipv4.protocol,
              hdr.ipv4.srcAddr,
              hdr.ipv4.dstAddr },
            hdr.ipv4.hdrChecksum,
            HashAlgorithm.csum16);
    }
}


/*************************************************************************
***********************  D E P A R S E R  *******************************
*************************************************************************/

control MyDeparser(packet_out packet, in headers hdr) {
    apply {
        /* TODO: add deparser logic */
        packet.emit(hdr.ethernet);
        packet.emit(hdr.arp);
        packet.emit(hdr.ipv4);
        packet.emit(hdr.ipv6);// 新增IPv6发送
        packet.emit(hdr.icmp);
        packet.emit(hdr.icmpv6);
    }
}

/*************************************************************************
***********************  S W I T C H  *******************************
*************************************************************************/

V1Switch(
MyParser(),
MyVerifyChecksum(),
MyIngress(),
MyEgress(),
MyComputeChecksum(),
MyDeparser()
) main;
